from copy import copy
from enum import Enum, auto
from itertools import count

from nanovllm.sampling_params import SamplingParams


class SequenceStatus(Enum):
    WAITING = auto()
    RUNNING = auto()
    FINISHED = auto()


class Sequence:
    """
    这段代码定义了 Sequence 类及其状态枚举 SequenceStatus，
    主要用于在大语言模型（LLM）推理过程中封装一个完整的生成任务（如一次对话、一段文本生成）的核心信息，
    包括生成的 token、状态、缓存块信息、采样参数等。它是调度器（Scheduler）管理任务的基本单元。
    """
    block_size = 256 # KV缓存块的大小（每个块可存储256个token的KV数据）
    counter = count() # 计数器，用于生成唯一的序列ID

    def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):
        self.seq_id = next(Sequence.counter) # 序列唯一ID（自增）
        self.status = SequenceStatus.WAITING # 初始状态为等待
        self.token_ids = copy(token_ids) # 生成的token列表（包含提示词+回答）
        self.last_token = token_ids[-1] # 最后一个token（用于后续生成）
        self.num_tokens = len(self.token_ids) # 总token数
        self.num_prompt_tokens = len(token_ids) # 提示词（输入）的token数
        self.num_cached_tokens = 0 # 已缓存到KV中的token数（避免重复计算）
        self.block_table = [] # 存储该序列占用的KV缓存块索引（与BlockManager配合）
        self.temperature = sampling_params.temperature
        self.max_tokens = sampling_params.max_tokens
        self.ignore_eos = sampling_params.ignore_eos

    def __len__(self):
        return self.num_tokens # 支持用len(seq)获取总token数

    def __getitem__(self, key):
        return self.token_ids[key]  # 支持用seq[key]获取指定位置的token（如seq[0]取第一个token）

    @property
    def is_finished(self):
        return self.status == SequenceStatus.FINISHED  # 判断序列是否已完成

    @property
    def num_completion_tokens(self):
        return self.num_tokens - self.num_prompt_tokens # 回答部分的token数（总token - 提示词token）

    @property
    def prompt_token_ids(self):
        return self.token_ids[:self.num_prompt_tokens]  # 提示词的token列表（输入部分）

    @property
    def completion_token_ids(self):
        return self.token_ids[self.num_prompt_tokens:] # 回答的token列表（生成部分）

    @property
    def num_cached_blocks(self):
        return self.num_cached_tokens // self.block_size  # 已缓存的块数量（缓存token数 / 块大小）

    @property
    def num_blocks(self):
        return (self.num_tokens + self.block_size - 1) // self.block_size # 总需要的块数量（向上取整：总token数 / 块大小）

    @property
    def last_block_num_tokens(self):
        # 若总 token 数为 300，block_size=256，则 num_blocks=2（256 + 44），last_block_num_tokens=44。
        return self.num_tokens - (self.num_blocks - 1) * self.block_size # 最后一个块中的token数（总token数 - 前面所有块的token数）

    def block(self, i):
        assert 0 <= i < self.num_blocks
        # 返回第i个缓存块中的token列表（按块大小切割）
        return self.token_ids[i*self.block_size: (i+1)*self.block_size]

    def append_token(self, token_id: int):
        self.token_ids.append(token_id)
        self.last_token = token_id
        self.num_tokens += 1

    def __getstate__(self):
        # 序列化时保存的状态：只保留必要信息，节省空间
        return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table,
                self.token_ids if self.num_completion_tokens == 0 else self.last_token)

    def __setstate__(self, state):
        # 反序列化时恢复状态
        self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table = state[:-1]
        if self.num_completion_tokens == 0:
            self.token_ids = state[-1] # 恢复完整token_ids
        else:
            self.last_token = state[-1]  # 只恢复最后一个token（节省空间）
